/* Copyright (C)  2005-2008 Infobright Inc.

This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License version 2.0 as
published by the Free  Software Foundation.

This program is distributed in the hope that  it will be useful, but
WITHOUT ANY WARRANTY; without even  the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
General Public License version 2.0 for more details.

You should have received a  copy of the GNU General Public License
version 2.0  along with this  program; if not, write to the Free
Software Foundation,  Inc., 59 Temple Place, Suite 330, Boston, MA
02111-1307 USA  */

#ifndef __COMPRESS_SUFFIXTREE_H
#define __COMPRESS_SUFFIXTREE_H

#include <vector>
#include <iostream>
#include <string.h>
#include "PPMdefs.h"

// this macro tells that some additional statistics should be computed and stored
// during suffix tree construction - which takes more time and memory, so finally should be switched off
#		if defined(_DEBUG) || (defined(__GNUC__) && !defined(NDEBUG))
#	define SUFTREE_STAT
#endif
//#define SUFTREE_CDAWG_SIZE

//------------------------------------------------------------------------

/*
	Ukkonen/McCreight's algorithm of suffix tree construction.
	The tree will need aprox. 20 times as much memory as the input data.
*/
template <class Symb = uchar, int NSymb = 256>
class SuffixTree : public PPMModel
{
public:
	// node identifier (pointer);
	// 0 - nil, 1 - root, positive - internal nodes, negative - leafs (-1 means the first leaf, i.e. number 0)
	typedef int PNode;

public:
	typedef uchar Bool;

	// internal node or leaf
	struct Node {
		int pos;		// int-1, leaf-2	// start of the edge label
		int len;		// int-2, leaf-2
		
		PNode suf;		// int-1, leaf-2	// suffix link; for PPM it's changed into _long_ suffix link
		
		int count;		// int-2, leaf-1
		Count sum;		// int-2, leaf-2	// upper bound of the node's range
		Bool cut;		// int-2, leaf-2	// true if the node was cut off during tree transformation for PPM
		
		Symb fsym;		// int-2, leaf-2
		PNode child;	// int-1, leaf-NO	// first child (NIL only for a leaf)
		PNode next;		// int-1, leaf-1	// sibling of this node (NIL for the last child)
		PNode prev;		// int-2, leaf-2	// previous sibling; cyclic list (first child points to the last)

#		ifdef SUFTREE_STAT
		int dep;
		int nvisit;			// how many times GetChild() was run for this node during tree construction
#		endif

		Node() { memset(this, 0, sizeof *this); }
	};

	// all nodes
	std::vector<Node> nodes;
	PNode NIL, ROOT;

	// original string - for reference
	Symb* data;
	int dlen;		// length of 'data', INcluding the terminating symbol ('\0')
	int* lens;		// lens[i] - length of substring of 'data' from 'i' to the next '\0' (incl. '\0')

	// global parameters and statistics
	int nleaves;
	bool iscreating;	// = true during tree construction, false otherwise
	PPMParam param;
	
	//-------------------------------------------------------------------------

	Node& GetNode(PNode n)          { return nodes[n]; }
	PNode NxtChild(PNode ch)		{ return GetNode(ch).next; }
	PNode PrvChild(PNode ch)		{ return GetNode(ch).prev; }
	PNode GetChild(PNode n, Symb s);
	void AddChild(Node& n, Symb s, PNode m);
	void AddChild(PNode n, Symb s, PNode m);
	void ChgChild(Node& n, Symb s, PNode m);

	//int GetLeafPos(PNode par, PNode leaf)	{ return GetNode(leaf).dep + GetNode(par).dep; }

	bool IsLeaf(PNode n)		{ return GetNode(n).child == NIL; }
	bool IsInternal(PNode n)	{ return GetNode(n).child != NIL; }

	struct Point { 
		PNode n;
		uint proj;		// number of symbols passed from 'n' down to 'next' (if Point at node, proj=0)
		PNode next;		// child of 'n' representing the current edge (if Point is at node, 'next' is undefined)
	};

	void Init();
	void Create();
	void Clear();

	// Moves the point 'p' along the string 's', as far as *s is the same as tree[p] (modifies 'p').
	// If the '\0' symbol was read, terminates and returns true. Otherwise returns false.
	// Assumption: at the beginning, 'p' is at a node (p.proj = 0).
	void MoveDown(Point& p, const uchar*& s);

	// Add new leaf below point 'p'; doesn't change p (even if a new node is created)
	void AddLeaf(Point p, int pos, PNode& last_int, PNode& last_leaf);

	// s - pointer to the first UNread character
	void MoveSuffix(Point& p, const uchar* s);
	void Canonize(Point& p, const uchar* s);

	//-------------------------------------------------------------------------
	// transformation of the tree for PPM
	
	// needed for derecursivation
	struct StackElem {
		PNode n;
		bool proc;		// are children of 'n' already processed?
		StackElem(PNode n_, bool proc_) { n=n_; proc=proc_; }
	};

	bool GetCut(PNode n)			{ return GetNode(n).cut == 1;   }
	void SetCut(PNode n, bool v)	{ GetNode(n).cut = (v ? 1 : 0); }
	Count GetEscCount(PNode n)		{ return param.esc_count;       }

	void SetCounts();
	
	// sets 'cut'=true for unnecessary nodes; physically breaks edges between pruned and non-pruned nodes	
	void Prune(PNode n, bool cut);
	bool IsValid(PNode n);			// is node 'n' useful for prediction? If not, its children will be cut off
	
	// 'sum' will be set only for not-cut-off nodes
	void SetSums();
	int Shift(int c);
	
	// after pruning some 'suf' links may be invalid and should be advanced to shorter suffixes;
	// advancing suf-links may be good also for compression performance
	void SetSuffix();
	bool Similar(PNode n, PNode suf);		// are distributions in 'n' and 'suf' similar? (suf is some suffix of n)

	//-------------------------------------------------------------------------
	// PPM compression/decompression
	
	PNode state;
	struct Edge {
		PNode n;	// e.n = NIL means "escape"
		Symb s;		// this is needed to identify exactly the edges from NIL to ROOT (otherwise is undefined)
	};

	void FindEdge(Edge& e, Symb* str, int len);		// compression
	void FindEdge(Edge& e, Count c);				// decompression
	
	CprsErr GetLabel(Edge e, Symb* lbl, int& len);		// 'len' - max size of lbl; upon exit: length of lbl
	int GetLen(Edge e);
	void GetRange(PNode stt, Edge e, Range& r);			// 'stt' must be the parent of 'e'
	Count GetTotal(PNode stt);

	void Move(Edge e);

//-----------------------------------------------------------------------------
public:
	// data_ - a sequence of concatenated null-terminated strings,
	// containing 'dlen_' symbols in total (INcluding the '\0' symbols!)
	// The resultant tree is a sum of trees constructed for each of the strings,
	// with every leaf storing the number of its repetitions.
	// If dlen_=-1, it's assumed that data_ contains only one string.
	SuffixTree(const Symb* data_, int dlen_ = -1)
	{
		Init();
		if(dlen_ < 0) dlen = (int)strlen((const char*)data_) + 1;
		else dlen = dlen_;
		dlen += 2;				// add '\0' at the beginning and at the end
		data = new uchar[dlen];
		memcpy(data+1, data_, dlen-2);
		data[0] = data[dlen-1] = '\0';
		Create();
	}
	virtual ~SuffixTree() { Clear(); }
	
	//-------------------------------------------------------------------------
	// information and statistics

	virtual int GetNNodes()   { return (int)nodes.size(); }
	virtual int GetNLeaves()  { return nleaves; }
	
	virtual int GetMemUsage();		// real number of bytes used, without wasted space in 'vector'
	virtual int GetMemAlloc();		// total number of bytes used

	// branch[256] - array of counts of branching factors
	// sumcnt[256] - array of summary subtree sizes (counts) for nodes of a given branching factor
	// sumvis[256] - array of summary number of visits (invoking GetChild) for a branching factor
	// vld - number of valid internal nodes, i.e. with large count and small branching - good for prediction
	//void GetStat(int* branch, int* sumcnt, int* sumvis, int& vld, int* cnt = 0, PNode n = 0, PNode prev = 0);
	
	// returns number of edges that would be removed if the tree were turned into Compact Directed Acyclic Word Graph
	int GetSizeCDAWG(PNode n = 1, bool cut = false);
	
	// flags (OR):
	//  1 - print label text
	//  2 - print lengths of labels
	//  4 - print depths of nodes
	//  8 - print counts
	void Print(std::ostream& str = std::cout, uint flags = 1+8, PNode n = 0, PNode prev = 0, int ind = 0);
	void PrintSub(std::ostream& str, int start, int stop = -1);		// stop=-1: print to '\0' symbol
	void PrintInd(std::ostream& str, int ind);
	
	//-------------------------------------------------------------------------
	// definitions for PPM

	virtual void TransformForPPM(PPMParam param_ = PPMParam());
	virtual void InitPPM()		{ state = ROOT; }

	// compression: [str,len_total] -> [len_of_edge,rng,total]
	virtual void Move(Symb* str, int& len, Range& rng, Count& total);

	// decompression: [c,str,len_max] -> [str,len_of_edge,rng]+returned_error
	virtual CprsErr Move(Count c, Symb* str, int& len, Range& rng);
	virtual Count GetTotal()	{ return GetTotal(state); }

#	ifdef SUFTREE_STAT
	int GetDep(PNode s)		{ return GetNode(s).dep; }
#	endif
};

//------------------------------------------------------------------------

template <class Symb, int NSymb>
 SuffixTree<>::PNode SuffixTree<Symb,NSymb>::GetChild(PNode n, Symb s) {
#	ifdef SUFTREE_STAT
	if(iscreating) GetNode(n).nvisit++;
#	endif
	if(n == NIL) return ROOT;
	PNode ch = GetNode(n).child;
	while(ch != NIL) 
		if(GetNode(ch).fsym == s) return ch;
		else ch = GetNode(ch).next;
	return NIL;
}

template <class Symb, int NSymb>
 void SuffixTree<Symb,NSymb>::AddChild(Node& n, Symb s, PNode m) 
{
	Node& mm = GetNode(m);
	mm.fsym = s;
	mm.next = n.child;
	if(n.child == NIL)
		mm.prev = m;
	else {
		PNode& prev = GetNode(n.child).prev;
		mm.prev = prev;
		prev = m;
	}
	n.child = m;
}
template <class Symb, int NSymb>
 void SuffixTree<Symb,NSymb>::AddChild(PNode n, Symb s, PNode m) 
{ 
	AddChild(GetNode(n), s, m); 
}

template <class Symb, int NSymb>
 void SuffixTree<Symb,NSymb>::ChgChild(Node& n, Symb s, PNode m) 
{
	// find a child with symbol 's'
	PNode* ch = &n.child;				// it's a pointer because we'll have to relink
	while(GetNode(*ch).fsym != s)		// the child must exist!!!
		ch = &GetNode(*ch).next;
	// replace the found child with node 'm'
	Node& mm = GetNode(m);
	mm.fsym = s;
	mm.next = GetNode(*ch).next;

	mm.prev = GetNode(*ch).prev;
	if(mm.prev == *ch)                     mm.prev = m;		// 'mm' is the only child
	else if(mm.next == NIL)  GetNode(n.child).prev = m;		// 'mm' is the last child out of 2 or more children
	else                     GetNode(mm.next).prev = m;		// 'mm' is not the last child

	*ch = m;
}

//------------------------------------------------------------------------

template <class Symb, int NSymb>
 int SuffixTree<Symb,NSymb>::GetLen(Edge e)
{
	if(e.n == NIL) return 0;
	return GetNode(e.n).len;
}

template <class Symb, int NSymb>
 Count SuffixTree<Symb,NSymb>::GetTotal(PNode stt)
{
	if(stt == NIL) return NSymb;
	Node& node = GetNode(stt);
	if(node.child == NIL) return param.esc_count;
	return GetNode(GetNode(node.child).prev).sum;
}

template <class Symb, int NSymb>
 void SuffixTree<Symb,NSymb>::Move(Edge e)
{ 
	if(e.n == NIL) 
		state = GetNode(state).suf; 
	else 
		state = e.n; 
}


#endif
